#encoding: utf-8

import re
from subprocess import Popen, PIPE
import os, tempfile, uuid

class OpenSSLException(Exception):
    def __init__(self, message):
        Exception.__init__(self)
        self.message = message

class TempFiles():
    """
    В качестве параметров конструктора принимает именованые параметры. Для каждого из параметров создается временный файл, в который записывается значение
    параметра, для использования возвращается массив вида {'имя параметра' : 'полный путь к временному файлу'}. При выходе из блока with файлы удаляются.
    
    К концу блока with все временные файлы должны быть закрыты, это достигается путем открытия файлов внутри блока так же через оператор with
    
    Пример использования:
    with TempFiles(asd='123') as files:
      with open(files['asd']) as file:
        print file.read()
    
    Напечатает 123
    """
    def __init__(self, **files):
        self.files_data = {}
        for file in files:
            self.files_data[file] = files[file]
            
    def __enter__(self):
        self.files_path = {}
        for file in self.files_data:
            fhandle, fpath = tempfile.mkstemp()
            if self.files_data[file] != None:
                os.write(fhandle, self.files_data[file])
            os.close(fhandle)
            self.files_path[file] = fpath
        return self.files_path
    
    def __exit__(self, type, value, traceback):
        for file in self.files_path:
            os.remove(self.files_path[file])

class OpenSSL:
    """
    Формат сообщения:
    keyname|sign (base64)|encrypted aes key (base64)|data (base64)
    """    
    def __init__(self, program, pubkey, privkey, trusted_folder):
        self.delim = '|'
        self.program = program
        self.pubkey = pubkey
        self.privkey = privkey
        self.trusted_folder = trusted_folder    
        self.priv_keyname = os.path.splitext(os.path.basename(privkey))[0]
        self.delim_size = len(self.delim)
    
    def _pubkey(self, keyname):
        if not re.match('^[a-zA-Z0-9\-_]+$', keyname):
            mess = "Incorrect keyname: '%s'" % keyname
            print mess
            raise OpenSSLException(mess)
        file = os.path.join(self.trusted_folder, keyname) + '.pub'
        return file
    
    def encrypt(self, data, keyname):
        """
        1. Генерирует ключ AES
        2. Шифрует данные ключем AES
        3. Шифрует ключ AES (base64) публичный ключем RSA
        4. Подписывает зашифрованые данные публичным ключем RSA
        5. Формурует сообщение

        keyname - имя публичного ключа получателя
        """
        
        # Сгенерировать ключ AES
        aes = uuid.uuid4().get_hex() + uuid.uuid4().get_hex()
        
        # Зашифровать данные
        with TempFiles(data=data, enc=None) as files:
            p = Popen([self.program, 'enc', '-aes-256-cbc', '-K', aes, '-iv', '0', '-in', files['data'], '-out', files['enc']])
            p.wait()
            with open(files['enc'], 'rb') as encfile:
                data = encfile.read()
                datalen = len(data)
        
        # Зашифровать ключ AES публичным ключем RSA
        p = Popen([self.program, 'rsautl', '-inkey', self._pubkey(keyname), '-pubin', '-encrypt'],
                  stdin=PIPE, stdout=PIPE, shell=False
                  )
        p.stdin.write(aes)
        p.stdin.close()
        p.wait()
        aes_enc = p.stdout.read()
        
        # Подпись
        with TempFiles(data=data) as files:
            p = Popen([self.program, 'dgst', '-sign', self.privkey, files['data']],
                      stdout=PIPE, shell=False
                      )
            p.wait()
            sign = p.stdout.read()
        
        res = self.priv_keyname + self.delim + sign.encode('base64') + self.delim + aes_enc.encode('base64') + self.delim + data.encode('base64')
        return res
    
    def decrypt(self, data):
        """
        Проверяет подпись по указаному ключу. Если подпись верна - расшифровывает, 
        иначе выкидывает исключение
        """
        keyname, sign, aes_enc, data_enc = data.split(self.delim, 4)
        
        sign, aes_enc, data_enc = sign.decode('base64'), aes_enc.decode('base64'), data_enc.decode('base64')
        
        # Проверка подписи
        with TempFiles(data_enc=data_enc, sign=sign) as files:
            p = Popen([self.program, 'dgst', '-signature', files['sign'], '-verify', self._pubkey(keyname), files['data_enc']],
                      stdout=PIPE, shell=False
                      )
            p.wait()
            result = p.stdout.read().strip()
            if result != 'Verified OK':
                raise OpenSSLException('Bad signature')
        
        # Расшифровать ключ AES
        with TempFiles(aes_enc=aes_enc,aes=None) as files:
            p = Popen([self.program, 'rsautl', '-inkey', self.privkey, '-decrypt', '-in', files['aes_enc'], '-out', files['aes']])
            p.wait()
            with open(files['aes'], 'rb') as fdata:
                aes = fdata.read()
                
        # Расшифровать данные
        with TempFiles(data_enc=data_enc, data=None) as files:
            p = Popen([self.program, 'enc', '-d', '-aes-256-cbc', '-K', aes, '-iv', '0', '-in', files['data_enc'], '-out', files['data']])
            p.wait()
            with open(files['data'], 'rb') as datafile:
                data = datafile.read()        
                
                
                
        return data, keyname